package com.sly.fuxi.kafka.config;

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;

import java.util.Map;

/**
 * @author SLY
 * @date 2025/3/3
 */
public class KafkaPartitioner implements Partitioner {

    @Override
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        // 获取分区总数
        int numPartitions = cluster.partitionCountForTopic(topic);
        // 获取指定的分区号
        if (key != null) {
            int specifiedPartition = Integer.parseInt(key.toString());
            // 确保指定的分区号在有效范围内
            if (specifiedPartition < 0 || specifiedPartition >= numPartitions) {
                throw new IllegalArgumentException("Invalid partition number: " + specifiedPartition);
            }
            return specifiedPartition;
        }
        return 0;
    }

    @Override
    public void close() {

    }

    @Override
    public void configure(Map<String, ?> configs) {

    }
}
